/*
 * ParaIdenCGRNN.cpp
 *
 *  Created on: Feb 6, 2015
 *      Author: phong
 */

#include "PDCGRNN.h"
#include "PDCorpus.h"
#include <vector>
#include <cmath>
#include "../Utils.h"
using namespace std;

real PDCGRNN::computeCostAndGrad(Container* cp, int startId, int endId, Param** gradList, unordered_set<int> &tWords, unordered_set<int> &tPath) {
	PDCorpus* tb = (PDCorpus*)cp;
	tWords.clear();
	tPath.clear();
	int nSample = 0;

	// prepare data, the first entry is the main
	for (int i = 0; i < nThreads; i++)
		gradList[i]->fill(0);

	real* costList = new real[nThreads]();
	int step = ceil((endId - startId+1) / (float)nThreads);

#pragma omp parallel for
	for (int th = 0; th < nThreads; th++) {
		for (int i = startId + th*step; i <= min(endId, startId + (th+1)*step-1); i++) {
			Deptree* cTree = std::get<1>(tb->storage[i]);
			Deptree* gTree = std::get<2>(tb->storage[i]);
			real score = std::get<0>(tb->storage[i]);
			compForward(cTree);
			real cost = genForward(gTree, cTree->cRep[0]);
			if (score > 0 || cost < 5000) {
				nSample++;
				costList[th] += score * cost;
				genBackprop(gTree, gradList[th], score);
				compBackprop(cTree, gradList[th], gTree->gradGRep[0]);
			}
			cTree->freeTempArrays();
			gTree->freeTempArrays();
		}
	}

	for (int i = startId; i <= endId; i++) {
		Deptree* cTree = std::get<1>(tb->storage[i]);
		Deptree* gTree = std::get<2>(tb->storage[i]);
		for (int k = 1; k < cTree->nNodes; k++) {
			int word = cTree->word[k];
			tWords.insert(word);
		}
		for (int k = 1; k < gTree->nNodes; k++) {
			int word = gTree->word[k];
			int* path = vocaDic->path[word];
			int codeLen = vocaDic->codeLength[word];
			for (int p = 0; p < codeLen; p++)
				tPath.insert(path[p]);
		}
	}

	Param* grad = gradList[0];

	// merge
	for (int i = 1; i < nThreads; i++) {
		for (int j = 0; j < grad->nMatricesWoWE; j++)
			grad->weights[j]->addi(gradList[i]->weights[j]);

		for (int p : tPath) {
			int j = p + grad->nMatricesWoWE;
			grad->weights[j]->addi(gradList[i]->weights[j]);
		}

		for (int w : tWords) {
			int j = w + grad->Wword.size() + grad->nMatricesWoWE;
			grad->weights[j]->addi(gradList[i]->weights[j]);
		}

		costList[0] += costList[i];
	}

	// compute regularities
	real cost = costList[0];
	cost /= nSample;

	for (int i = 0; i < params->nMatricesWoWE; i++) {
		Matrix* W = params->weights[i];
		cost += Utils::sumSqr(W->data, W->length) * lambda/2;
		grad->weights[i]->divi(nSample)->addi(lambda, W);
	}

	for (int word : tWords) {
		Matrix* wemb = params->L[word];
		cost += Utils::sumSqr(wemb->data, wemb->length) * lambdaL/2;
		grad->L[word]->divi(nSample)->addi(lambdaL, wemb);
	}

	for (int p : tPath) {
		Matrix* v = params->Wword[p];
		cost += Utils::sumSqr(v->data, v->length) * lambda/2;
		grad->Wword[p]->divi(nSample)->addi(lambda, v);
	}

	delete[] costList;

	return cost;
}
